package passlock

import scala.sys.process._
import chisel3._
import chisel3.util._
import Config._
import chisel3.experimental.chiselName

object Config{
  val UnlockFSMwidth = 2
  val SuperFSMwidth = 3
  val ResetPasswordFSMwidth = 3
  val DataWidth = 4
  val Passwordwidth = 2
  val Inputwidth = 12
  val Passwordlength: Int = 1 << Passwordwidth
}

@chiselName
class UnlockFSM extends Module{
  val io = IO(new Bundle {
    val valid: Bool = Input(Bool())
    val R1: Bool = Input(Bool())
    val R2: Bool = Input(Bool())
    val R3: Bool = Input(Bool())
    val R4: Bool = Input(Bool())
    val R13: Bool = Input(Bool())
    val Unlock: Bool = Output(Bool())
    val Unlocking: Bool = Input(Bool())
  })
  withReset(reset.asBool() || io.Unlocking){
    val state: UInt = RegInit(0.U(UnlockFSMwidth.W))
    state := MuxCase(0.U, Array(
      !io.valid -> state,
      (state === 0.U(UnlockFSMwidth.W)) -> MuxCase(0.U, Array(
        io.R1 -> 1.U(UnlockFSMwidth.W),
        (!io.R1) -> 0.U(UnlockFSMwidth.W)
      )),
      (state === 1.U(UnlockFSMwidth.W)) -> MuxCase(0.U, Array(
        io.R2 -> 3.U(UnlockFSMwidth.W),
        io.R1 -> 1.U(UnlockFSMwidth.W)
      )),
      (state === 3.U(UnlockFSMwidth.W)) -> MuxCase(0.U, Array(
        io.R3 -> 2.U(UnlockFSMwidth.W),
        (io.R1 && io.R2) -> 3.U(UnlockFSMwidth.W),
        (io.R1 && !io.R2) -> 1.U(UnlockFSMwidth.W)
      )),
      (state === 2.U(UnlockFSMwidth.W)) -> MuxCase(0.U, Array(
        io.R4 -> 0.U(UnlockFSMwidth.W),
        (io.R1 && io.R2 && io.R13) -> 2.U(UnlockFSMwidth.W),
        (!io.R1 && io.R2 && io.R13) -> 3.U(UnlockFSMwidth.W),
        (io.R1 && (!io.R13 || !io.R2)) -> 1.U(UnlockFSMwidth.W)
      ))
    ))
    io.Unlock := RegNext(Mux(io.valid, (state === 2.U(UnlockFSMwidth.W)) && io.R4, io.Unlock), false.B)
  }
}

class SuperFSM extends Module{
  val io = IO(new Bundle {
    val DataInput: UInt = Input(UInt(DataWidth.W))
    val Activate: Bool = Output(Bool())
    val valid: Bool = Input(Bool())
  })
  val state: UInt = RegInit(0.U(SuperFSMwidth.W))
  state := MuxCase(0.U, Array(
    !io.valid -> state,
    (state === 0.U(SuperFSMwidth.W)) -> Mux(io.DataInput === 2.U, 1.U(SuperFSMwidth.W), 0.U(SuperFSMwidth.W)),
    (state === 1.U(SuperFSMwidth.W)) -> Mux(io.DataInput === 9.U, 2.U(SuperFSMwidth.W), Mux(io.DataInput === 1.U, 1.U(SuperFSMwidth.W), 0.U(SuperFSMwidth.W))),
    (state === 2.U(SuperFSMwidth.W)) -> Mux(io.DataInput === 1.U, 3.U(SuperFSMwidth.W), Mux(io.DataInput === 1.U, 1.U(SuperFSMwidth.W), 0.U(SuperFSMwidth.W))),
    (state === 3.U(SuperFSMwidth.W)) -> Mux(io.DataInput === 5.U, 4.U(SuperFSMwidth.W), Mux(io.DataInput === 1.U, 1.U(SuperFSMwidth.W), 0.U(SuperFSMwidth.W))),
    (state === 4.U(SuperFSMwidth.W)) -> Mux(io.DataInput === 1.U, 5.U(SuperFSMwidth.W), Mux(io.DataInput === 1.U, 1.U(SuperFSMwidth.W), 0.U(SuperFSMwidth.W))),
    (state === 5.U(SuperFSMwidth.W)) -> Mux(io.DataInput === 10.U, 0.U(SuperFSMwidth.W), Mux(io.DataInput === 1.U, 1.U(SuperFSMwidth.W), 0.U(SuperFSMwidth.W)))
  ))
  io.Activate := RegNext(Mux(io.valid, (state === 5.U(SuperFSMwidth.W)) && (io.DataInput === 10.U), io.Activate), false.B)
}

class ResetPasswordFSM extends Module{
  val io = IO(new Bundle {
    val Activate: Bool = Input(Bool())
    val valid: Bool = Input(Bool())
    val SaveCtr: Bool = Output(Bool())
    val CompareCtr: UInt = Output(UInt(Passwordwidth.W))
    val DataInCtr: Bool = Output(Bool())
    val SymbolStar: Bool = Input(Bool())
  })
  val state: UInt = RegInit(0.U(ResetPasswordFSMwidth.W))
  state := Mux(io.valid, Mux(((state === 0.U) && !io.Activate) || io.SymbolStar, 0.U, state + 1.U), state)
  io.SaveCtr := RegNext(Mux(io.valid, state === 7.U(ResetPasswordFSMwidth.W) && !io.SymbolStar, io.SaveCtr), false.B)
  io.CompareCtr := state(1, 0)
  io.DataInCtr := !state(2)
}

class ResetPassword extends Module{
  val io = IO(new Bundle {
    val Password1: UInt = Output(UInt(DataWidth.W))
    val Password2: UInt = Output(UInt(DataWidth.W))
    val Password3: UInt = Output(UInt(DataWidth.W))
    val Password4: UInt = Output(UInt(DataWidth.W))
    val ResetPassword: Bool = Output(Bool())
    val DataIn: UInt = Input(UInt(DataWidth.W))
    val SaveCtr: Bool = Input(Bool())
    val CompareCtr: UInt = Input(UInt(Passwordwidth.W))
    val DataInCtr: Bool = Input(Bool())
    val Shape: Bool = Input(Bool())
    val valid: Bool = Input(Bool())
    val OuterReset: Bool = Input(Bool())
  })
  val DataReg: Vec[UInt] = RegInit(VecInit(Seq.fill(Passwordlength)(0.U(DataWidth.W))))
  val StoreCtrReg: Bool = RegInit(false.B)
  StoreCtrReg := Mux(io.valid, io.DataInCtr || StoreCtrReg && (io.DataIn === MuxCase(0.U, DataReg.zipWithIndex.map(tup => (io.CompareCtr ===(3.U - tup._2.U), tup._1)))), StoreCtrReg)
  withReset(io.OuterReset){
    val StoreReg: Vec[UInt] = RegInit(VecInit(Seq.fill(Passwordlength)(1.U(DataWidth.W))))
    StoreReg.zip(DataReg).foreach(tup => tup._1 := Mux(io.valid, Mux(io.ResetPassword && io.Shape, tup._2, tup._1), tup._1))
    io.Password1 := StoreReg(3)
    io.Password2 := StoreReg(2)
    io.Password3 := StoreReg(1)
    io.Password4 := StoreReg(0)
  }
  val LastReg: UInt = DataReg.foldLeft(io.DataIn){ (left, right) => {
    right := Mux(io.DataInCtr && io.valid, left, right)
    right
  }}
  io.ResetPassword := StoreCtrReg && io.SaveCtr
}

@chiselName
class Core extends Module{
  val io = IO(new Bundle {
    val DataIn: UInt = Input(UInt(Inputwidth.W))
    val Unlock: Bool = Output(Bool())
    val Unlocking: Bool = Input(Bool())
    val Halt: Bool = Output(Bool())
    val Halting: Bool = Input(Bool())
    val SymbolStar: Bool = Output(Bool())
  })
  val OuterReset: Bool = reset.asBool()
  val SymbolStar: Bool = !io.DataIn(Inputwidth - 1)
  val SymbolShape: Bool = !io.DataIn(Inputwidth - 2)
  val DataIn: UInt = OHToUInt((~io.DataIn)<<1)
  val valid: Bool = RegNext(DataIn) =/= DataIn && DataIn=/= 0.U
  io.SymbolStar := SymbolStar
  withReset(reset.asBool() || io.Halting){
    val MResetPassword: ResetPassword = Module(new ResetPassword)
    val MUnlockFSM: UnlockFSM = Module(new UnlockFSM)
    val MSuperFSM: SuperFSM = Module(new SuperFSM)
    val MResetPasswordFSM: ResetPasswordFSM = Module(new ResetPasswordFSM)
    val FalsePassword = !MUnlockFSM.io.Unlock && !MResetPassword.io.ResetPassword && !io.Unlocking && SymbolShape
    MUnlockFSM.io.R1 := DataIn === MResetPassword.io.Password1
    MUnlockFSM.io.R2 := DataIn === MResetPassword.io.Password2
    MUnlockFSM.io.R3 := DataIn === MResetPassword.io.Password3
    MUnlockFSM.io.R4 := DataIn === MResetPassword.io.Password4
    MUnlockFSM.io.R13 := MResetPassword.io.Password1 === MResetPassword.io.Password3
    MUnlockFSM.io.Unlocking := io.Unlocking
    MUnlockFSM.io.valid := valid
    MSuperFSM.io.DataInput := DataIn
    MSuperFSM.io.valid := valid
    MResetPasswordFSM.io.Activate := MSuperFSM.io.Activate
    MResetPasswordFSM.io.valid := valid
    MResetPasswordFSM.io.SymbolStar := io.SymbolStar
    MResetPassword.io.SaveCtr := MResetPasswordFSM.io.SaveCtr
    MResetPassword.io.DataInCtr := MResetPasswordFSM.io.DataInCtr
    MResetPassword.io.CompareCtr := MResetPasswordFSM.io.CompareCtr
    MResetPassword.io.DataIn := DataIn
    MResetPassword.io.Shape := SymbolShape
    MResetPassword.io.valid := valid
    MResetPassword.io.OuterReset := OuterReset
    io.Unlock := MUnlockFSM.io.Unlock && SymbolShape && !MResetPassword.io.ResetPassword
    withReset(OuterReset || io.Unlock) {
      val HaltState: UInt = RegInit(0.U(2.W))
      HaltState := MuxCase(0.U, Array(
        (!FalsePassword || !valid || io.Halting) -> HaltState,
        (HaltState === 0.U) -> 1.U,
        (HaltState === 1.U) -> 3.U,
        (HaltState === 3.U) -> 0.U
      ))
      io.Halt := HaltState === 3.U && FalsePassword && valid
    }
  }
}

object Top {
  def main(args: Array[String]): Unit = {
    chisel3.Driver.execute(args, () => new Core())
    "dot -Tpng ./src/main/dot/SuperFSM.dot -o ./SuperFSM.png".!
    "dot -Tpng ./src/main/dot/Unlock.dot -o ./Unlock.png".!
    "dot -Tpng ./src/main/dot/ResetPasswordFSM.dot -o ./ResetPasswordFSM.png".!
    "iverilog -o Sim Sim.v Core.v".!
    "vvp -n Sim -lxt2".!
    "cp Sim.vcd Sim.lxt".!
    "rm Sim.vcd".!
    "rm Sim".!
    "gtkwave Sim.lxt".!
    "rm Sim.lxt".!
  }
}